import pygame
from snake import Snake
from map import GameMap
from food import Bait,Bonus
import random
import sys
from constants import * 
import numpy as np 
from pygame import mixer , time

class Game:
    maps = None
    
    def __init__(self,silent_mode = True):
        
        if Game.maps is None:
            Game.maps = GameMap.load_maps()
        self.current_map_index = 0
        self.map = Game.maps[self.current_map_index]
        # 游戏画面的大小
        self.width_cell_num = self.map.width  # 宽度上的格子数
        self.height_cell_num = self.map.height   # 高度上的格子数
        self.grid_num = self.width_cell_num * self.height_cell_num  # 总共的格子数
        self.cell_size = self.map.cell_size
        self.width = self.width_cell_num * self.cell_size
        self.height = self.height_cell_num * self.cell_size
        
        
        # 整个界面的大小，囊括了游戏画面
        self.border_length = 20 # 边框长度
        self.display_width = self.width + 2 * self.border_length 
        self.display_height = self.height + 2 * self.border_length + self.cell_size # 在高度上多一个cell的大小 

        #### 
        self.score = 0
        self.level = 1
        ### 
        self.silent_mode = silent_mode
        if not silent_mode : 
            # 加载音效
            mixer.init() 
            self.sound_eat = mixer.Sound('sound/eat.wav')
            self.count_down = mixer.Sound('sound/eat.wav')
            self.sound_game_over = mixer.Sound('sound/game_over.wav')
            self.sound_victory = mixer.Sound('sound/victory.wav')
    
        self.snake = None
        self.non_snake = None
        self.bait = None
        self.bonus = None
        self.bonus_clock = None
        
        
        # these four lines code is no longer needed 
        # train 的时候注释掉，test的时候打开
        # pygame.init()
        # pygame.display.set_caption("Greedy Snake Game")
        # self.screen = pygame.display.set_mode((self.display_width, self.display_height))
        # self.font = pygame.font.Font(None, 36)
        

        self.reset() 
        
    def reset(self) :
        # 初始的蛇有三个cell，三个cell以元组的形式存放在列表里
        self.snake = Snake(self.height_cell_num//2, self.width_cell_num//2)
        # 初始化不被蛇占据的棋盘
        self.non_snake = set([(row,col) for row in range(self.height_cell_num) for col in range(self.width_cell_num) if (row,col) not in self.snake.positions])
        self.bait = self._generate_food("bait")
        self.bonus = None
        self.bonus_clock = 0 
        self.score = 0 
    
    def _generate_food(self , type : str):
        if type == 'bait' :
            x, y = self.map.get_random_position(self.snake.positions)
            return Bait(x, y)
        elif type == 'bonus' :
            x , y = 0,0 
            while True :
                x, y = self.map.get_random_position(self.snake.positions)
                if (x,y) not in self.bait.position:
                    return Bonus(x, y)

    def step(self,action):
        die , x , y  = self.snake.move(action,self.map)

        if not die :
            self.snake.positions.insert(0,(x,y)) # 将蛇头加入到列表首位
            self.non_snake.remove((x,y)) # 蛇头位置占据了棋盘某格
        else :
            # game over
            if len(self.snake.positions) < self.grid_num : # 没有完全胜利
                if not self.silent_mode :
                    self.sound_game_over.play() 
            else :
                if not self.silent_mode  :
                    # 完全胜利
                    self.sound_victory.play() 
                    
        food_obtained = False
        if (x, y) == self.bait.position :
            food_obtained = True
            self.score += 10 
            # 非静音模式
            if not self.silent_mode :
                self.sound_eat.play() 
        elif self.bonus and (x,y) == self.bonus.position :
            self.score += 20 * (time.get_ticks() - self.bonus_clock ) // 1000 
            # 非静音模式
            if not self.silent_mode :
                self.sound_eat.play() 
            self.bonus = None  # 删除bonus
        else :
            food_obtained = False
            # 没吃到食物，剔除蛇列表的末尾加入到非蛇区
            self.non_snake.add(self.snake.positions.pop())
        

        # 如果吃到了bait,按概率刷新bonus
        if food_obtained :
            self.bait = self._generate_food("bait") 
            if not self.bonus and  random.random() < 0.1 : # 如果没有bonus并且抽中了bonus，就刷新出来
                self.bonus = self._generate_food("bonus")
                self.bonus_clock = time.get_ticks()
            
        if self.bonus and (time.get_ticks() - self.bonus_clock ) // 1000 >  10 :
            self.bonus = None
            self.bonus_clock = None
            
        
        info = {
            "snake_length" : len(self.snake) ,
            "snake_head_position" : np.array(self.snake[0]) ,
            "prev_snake_head_position" : np.array(self.snake[1]) ,
            "bait_pos" : np.array(self.bait.position) ,
            "bonus_pos" : np.array(self.bonus.position) if self.bonus is not None else None , 
            "food_obtained" : food_obtained 
        }
        return die , info 
    
    
    def render(self):
        self.screen.fill((0,0,0))
        
        # 绘制障碍
        self.draw_obstacles(self.map.obstacles)
        
        # 绘制墙壁
        self.draw_walls() 

        # 绘制蛇
        self.draw_snake() 
        
        # 绘制分数和等级界面
        self.show_score_and_level()
        
        self.draw_food()

        pygame.display.flip() 
        # 只处理退出事件
        for e in pygame.event.get() :
            if e.type == pygame.QUIT :
                pygame.quit() 
                sys.exit()
    
    def draw_food(self):
        # 如果蛇没有占满全屏，绘制食物
        if len(self.snake) < self.grid_num :
            row,col = self.bait.position 
            # lure
            pygame.draw.rect(self.screen , Bait.COLOR , 
            (col * self.cell_size +  self.border_length , row * self.cell_size + self.border_length , self.cell_size , self.cell_size))

            # bonus 
            if self.bonus :
                row,col = self.bonus.position 
                # bonus 
                pygame.draw.rect(self.screen , Bonus.COLOR , 
                (col * self.cell_size +  self.border_length , row * self.cell_size + self.border_length , self.cell_size , self.cell_size))
                

    
    def draw_tasty_food(self):
        # 如果蛇没有占满全屏，绘制食物
        if len(self.snake) < self.grid_num :
            row,col = self.food 
            # 红色
            pygame.draw.rect(self.screen , (255,0,0) , (col * self.cell_size +  self.border_length , row * self.cell_size + self.border_length , self.cell_size , self.cell_size))
    
    
    def draw_obstacles(self,obstacles):
        for row , col in obstacles : #
            body_x = col * self.cell_size + self.border_length
            body_y = row * self.cell_size + self.border_length
            body_width  , body_height = self.cell_size , self.cell_size 
            body_radius = 5 
            pygame.draw.rect(self.screen , WHITE,
                            (body_x , body_y , body_width , body_height) , 
                            border_radius=body_radius)
        
    
    def draw_walls(self):
        # 绘制边界
        pygame.draw.rect(self.screen , WHITE , (self.border_length  , self.border_length , self.width , self.height ) , 5 ) 


    def draw_welcome_screen(self) :
        '''
        游戏开始界面绘制
        '''
        title_text = self.font.render('SNAKE GAME' , True , WHITE) 
        start_button_text = "START" 
        
        # 黑色填充
        self.screen.fill((0,0,0))
        self.screen.blit(title_text,(self.display_width // 2 - title_text.get_width() // 2 , self.display_height // 4 ))
        self.draw_button_text(start_button_text,(self.display_width // 2 , self.display_height // 2))
        # 更新整个 Pygame 显示窗口。所有绘图完成之后，将缓冲区中的内容显示到屏幕上，并清空缓冲区以便下一帧的绘制
        pygame.display.flip() 
        
        
    def draw_game_over_screen(self):
        # surface "GAME OVER"字样
        game_over_text = self.font.render("GAME OVER",True , WHITE)
        final_score_text = self.font.render(f"SCORE:{self.score}" , True , (255,255,255)) 
        retry_button_text = "RETRY"
        
        self.screen.fill((0,0,0))
        self.screen.blit(game_over_text,(self.display_width // 2 - game_over_text.get_width() // 2 , self.display_height // 4 ))
        self.screen.blit(final_score_text , (self.display_width // 2 - final_score_text.get_width() // 2 , self.display_height // 4 + final_score_text.get_height() + 10 ))
        self.draw_button_text(retry_button_text , (self.display_width // 2 , self.display_height // 2 )) 
        pygame.display.flip() 
        
        
    def draw_button_text(self , button_text_str , pos , hover_color=WHITE , normal_color = (100,100,100)) :
        mouse_pos = pygame.mouse.get_pos() 
        button_text = self.font.render(button_text_str , True , normal_color) 
        text_rect = button_text.get_rect(center=pos)
        
        if text_rect.collidepoint(mouse_pos):
            colored_text = self.font.render(button_text_str , True , hover_color) 
        else :
            colored_text = self.font.render(button_text_str , True , normal_color) 
        
        self.screen.blit(colored_text , text_rect) 
        
    
    def draw_countdown(self,number) :
        countdown_text = self.font.render(str(number) , True , WHITE)
        self.screen.blit(countdown_text,(self.width // 2 - countdown_text.get_width() // 2 , self.height // 2 - countdown_text.get_height() // 2 ))
        pygame.display.flip() 
        
        
    def is_mouse_on_button(self,button_text):
        mouse_pos = pygame.mouse.get_pos() 
        text_rect = button_text.get_rect(
            center = (
                self.display_width // 2 ,
                self.display_height // 2  ,
            )
        )
        return text_rect.collidepoint(mouse_pos) 
    

    def draw_snake(self):
        # 获得蛇头的行列,并确定绘制位置
        head_row , head_col = self.snake.positions[0] 
        head_x  = head_col * self.cell_size + self.border_length
        head_y  = head_row * self.cell_size + self.border_length

         
        # 蛇头始终是一个菱形
        pygame.draw.polygon(self.screen , (100,100,255) , [
            (head_x + self.cell_size // 2, head_y),
            (head_x + self.cell_size, head_y + self.cell_size // 2),
            (head_x + self.cell_size // 2, head_y + self.cell_size),
            (head_x, head_y + self.cell_size // 2)
        ])
        
        # 点缀蛇眼
        eye_size = 3
        eye_offset = self.cell_size // 4 
        # 绘制左眼
        pygame.draw.circle(self.screen , WHITE , (head_x+ eye_offset , head_y + eye_offset) , eye_size)
        # 绘制右眼
        pygame.draw.circle(self.screen , WHITE , (head_x + self.cell_size - eye_offset , head_y + eye_offset) , eye_size) 
        
        # 绘制蛇身
        # 渐变色
        color_list = np.linspace(255,100,len(self.snake),dtype=np.uint8)
        i = 1 
        body_width  , body_height = self.cell_size , self.cell_size
        body_radius = 5  
        body_x , body_y = 0,0
        for row , col in self.snake[1:] : # 除了蛇头其余依次绘制
            body_x = col * self.cell_size + self.border_length
            body_y = row * self.cell_size + self.border_length
            
            pygame.draw.rect(self.screen , (0 , color_list[i] , 0),
                            (body_x , body_y , body_width , body_height) , 
                            border_radius=body_radius)
    
            i += 1 
        # 重绘最后一节身体
        pygame.draw.rect(self.screen , (255,100,100) , 
                        (body_x , body_y , body_width , body_height) , 
                        border_radius = body_radius)

    
    
    def show_score_and_level(self):
        font = pygame.font.Font(None, 36)
        text = font.render(f"Score: {self.score}", True, WHITE)
        self.screen.blit(text,(self.border_length,self.height + 2 * self.border_length))
        text = font.render(f"Level: {self.level}", True, WHITE)
        self.screen.blit(text,(self.width - text.get_width()  ,self.height + 2 * self.border_length))